import os
import argparse
import joblib
import csv
from sample_predictor import predict_sample

# 配置文件路径
MODEL_PATH = "models/codebert-base"
LABEL_FILE = "samples/selectedLabels.csv"
CUSTOM_FILE = "customs/codebert_selectedpro_classifier.joblib"
SAMPLE_PATH = "samples/subtrain/selectedpro/test/1eJx34l8pcAFvMuOwrjB.asm"  # 测试样本路径

def load_labels(label_file):
    """
    加载标签数据
    """
    labels = {}
    if os.path.exists(label_file):
        with open(label_file, 'r') as f:
            reader = csv.DictReader(f)
            for row in reader:
                labels[row['Id']] = int(row['Class'])
    return labels

def get_file_id(file_path):
    """
    从文件路径中提取文件ID
    """
    filename = os.path.basename(file_path)
    return os.path.splitext(filename)[0]

def main():
    # 创建命令行参数解析器
    parser = argparse.ArgumentParser(description='木马病毒分类器 - 单文件预测')
    parser.add_argument('--file', type=str, default=SAMPLE_PATH,
                        help=f'要预测的ASM文件路径（默认: {SAMPLE_PATH}）')
    parser.add_argument('--model_path', type=str, default=MODEL_PATH,
                        help=f'BERT模型路径 (默认: {MODEL_PATH})')
    parser.add_argument('--model_file', type=str, default=CUSTOM_FILE,
                        help=f'分类模型文件路径 (默认: {CUSTOM_FILE})')
    parser.add_argument('--label_file', type=str, default=LABEL_FILE,
                        help=f'标签文件路径 (默认: {LABEL_FILE})')
    
    args = parser.parse_args()
    
    # 检查文件是否存在
    if not os.path.exists(args.file):
        print(f"错误: 文件 {args.file} 不存在!")
        return
    
    # 加载分类模型
    if not os.path.exists(args.model_file):
        print(f"错误: 模型文件 {args.model_file} 不存在!")
        print("请先运行 main.py 训练模型")
        return
    
    # 加载标签数据
    labels = load_labels(args.label_file)
    
    # 获取文件ID
    file_id = get_file_id(args.file)
    
    # 加载模型
    pipeline = joblib.load(args.model_file)
    
    # 使用已实现的predict_sample接口进行预测
    prediction, probabilities = predict_sample(
        args.model_path, 
        args.file, 
        pipeline
    )
    
    # 输出结果
    if prediction is not None:
        print("\n" + "="*50)
        print(f"文件: {args.file}")
        print(f"文件ID: {file_id}")
        
        # 如果标签文件中包含该文件ID，显示真实标签
        if file_id in labels:
            true_label = labels[file_id]
            print(f"真实类别: {true_label}")
            print(f"预测类别: {prediction}")
            print(f"预测结果: {'正确' if prediction == true_label else '错误'}")
        else:
            print(f"预测类别: {prediction}")
            print("注意: 标签文件中未找到该文件ID")
        
        print("\n各类别概率分布:")
        for i, prob in enumerate(probabilities, start=1):
            print(f"  类别 {i}: {prob*100:.2f}%")
        print("="*50)
    else:
        print("预测失败，请检查文件格式或模型")

if __name__ == "__main__":
    main()